# +----------------------------------------------------------------------
# | ChatWork
# +----------------------------------------------------------------------
# | 这不是一个自由软件,您只能在不用于商业目的的前提下对程序代码进行修改和使用。
# | 任何企业和个人不允许对程序代码以任何形式任何目的再发布,商业使用请获取授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import time
from typing import List
from tortoise.queryset import Q
from tortoise.transactions import in_transaction
from hypertext import PagingResult
from exception import AppException
from common.utils.urls import UrlUtil
from common.utils.times import TimeUtil
from common.utils.tools import ToolsUtil
from common.enums.agent import AgentEnum
from common.enums.models import ModelsEnum
from common.models.users import UserModel
from common.models.dev import DevModelsModel
from common.models.agents import AgentRobotModel
from common.models.agents import AgentKnowModel
from common.models.agents import AgentKnowTeamedModel
from common.postgres.public import PgKnowledgeModel
from common.service.clean_logic import CleanLogic
from common.service.subscribe_logic import SubscribeLogic
from plugins.censor.driver import CensorDriver
from apps.api.schemas.agent.know import manage_schema as schema
from apps.api.service.agent.know.teamed_service import KnowTeamedService
from apps.api.service.agent.robot.teamed_service import RobotTeamedService


class KnowManageService:
    """ 知识库服务类 """

    @classmethod
    async def whole(cls, user_id: int, params: schema.KnowWholeIn) -> List[schema.KnowWholeVo]:
        """
        所有知识库。

        Args:
            user_id (int): 用户ID。
            params (schema.KnowWholeIn): 搜索参数。

        Returns:
            List[schema.KnowWholeVo]: 知识库列表Vo。

        Author:
            zero
        """
        uid: int = user_id

        # 智能体过滤
        if params.bot:
            # 查智能体
            robot = await AgentRobotModel.filter(code=params.bot).get()
            # 拥有者ID
            uid = robot.user_id
            # 验证权限
            await RobotTeamedService.has_permission(
                bot=params.bot,
                user_id=user_id,
                power=RobotTeamedService.POWER_VIEW
            )

        # 查询知识库
        lists = await (AgentKnowModel
                       .filter(user_id=uid)
                       .filter(is_delete=0)
                       .order_by("-sort", "-update_time")
                       .all())

        # 查询模型
        models = []
        vector = [item.vector_model.split(":")[0] for item in lists if item.vector_model]
        if vector:
            models = await (DevModelsModel
                            .filter(mark__in=list(set(vector)))
                            .filter(scene=ModelsEnum.TYPE_VECTOR)
                            .filter(is_delete=0)
                            .all().values_list("mark", "name"))
            models = {k: v for k, v in models}

        # 处理数据
        _lists = []
        for item in lists:
            mark = item.vector_model.split(":")[0]
            _lists.append(schema.KnowWholeVo(
                code=item.code,
                logo=await UrlUtil.to_absolute_url(item.logo),
                name=item.name,
                vector_name=models.get(mark) or "-",
                vector_model=item.vector_model,
                is_disable=item.is_disable,
            ))
        return _lists

    @classmethod
    async def lists(cls, params: schema.KnowSearchIn, user_id: int) -> PagingResult[schema.KnowListVo]:
        """
        知识库列表。

        Args:
            params (schema.KnowSearchIn): 知识库查询参数。
            user_id (int): 用户ID。

        Returns:
            PagingResult[schema.KnowListVo]: 知识库分页列表Vo。

        Author:
            zero
        """
        # 构建查询条件
        where = [Q(is_delete=0)]
        if params.keyword:
            where.append(Q(name__icontains=params.keyword))

        # 成员身份过滤
        permissions = {user_id: AgentEnum.POWER_OWNED}
        if params.owner == "private":
            where.append(Q(user_id=user_id))
        elif params.owner == "cooperate":
            team_lis = await AgentKnowTeamedModel.filter(user_id=user_id).all()
            team_ids = [item.know_id for item in team_lis if item.know_id]
            permissions = {team.user_id: team.power for team in team_lis}
            where.append(Q(id__in=team_ids or [0]))
        else:
            team_lis = await AgentKnowTeamedModel.filter(user_id=user_id).all()
            team_ids = [int(item.know_id) for item in team_lis if item.know_id]
            if team_ids:
                permissions = {team.user_id: team.power for team in team_lis}
                where.append(Q(user_id=user_id) | Q(id__in=team_ids or [0]))
            else:
                where.append(Q(user_id=user_id))

        # 查询知识库列表
        _model = AgentKnowModel.filter(*where).order_by("-sort", "-update_time")
        _pager = await AgentKnowModel.paginate(
            model=_model,
            page_no=params.page_no,
            page_size=params.page_size
        )

        # 查询用户信息
        user_ids = [item["user_id"] for item in _pager.lists if item["user_id"]]
        users = await UserModel.queryUserDictsById(user_ids)

        # 查询模型信息
        models = [item["vector_model"] for item in _pager.lists if item["vector_model"]]
        vector = await DevModelsModel.get_model_alias(models, ModelsEnum.TYPE_VECTOR)

        # 循环处理列表
        _lists = []
        for item in _pager.lists:
            logo: str = await UrlUtil.to_absolute_url(item["logo"])
            owner: bool = True if item["user_id"] == user_id else False
            power: int = permissions.get(item["user_id"], 0)
            embed: dict = vector.get(item["vector_model"], {})
            user: dict = users.get(item["user_id"], {})

            _lists.append(schema.KnowListVo(
                code=item["code"],
                logo=logo,
                name=item["name"],
                intro=item["intro"],
                power=power,
                is_owner=owner,
                vector_model=embed.get("model"),
                is_disable=item["is_disable"],
                create_time=item["create_time"],
                update_time=item["update_time"],
                user={
                    "sn": user.get("sn", ""),
                    "nickname": user.get("nickname", ""),
                    "avatar": await UrlUtil.to_absolute_url(user.get('avatar', ""))
                }
            ))

        _pager.lists = _lists
        return _pager

    @classmethod
    async def detail(cls, code: str, user_id: int) -> schema.KnowDetailVo:
        """
        知识库详情。

        Args:
            code (str): 知识库编号。
            user_id (int): 用户ID。

        Returns:
            schema.KnowDetailVo: 知识库详情Vo。

        Author:
            zero
        """
        # 查询知识库
        know = await AgentKnowModel.filter(code=code, is_delete=0).get()

        # 权限的验证
        if know.user_id != user_id:
            await KnowTeamedService.has_permission(know.code, user_id, AgentEnum.POWER_VIEW)

        # 查询模型
        marks = know.vector_model.split(":")[0]
        model = await DevModelsModel.filter(scene=ModelsEnum.TYPE_VECTOR, mark=marks).first()
        vector_name: str = model.name if model else "-"

        # 权限查询
        power: int = KnowTeamedService.POWER_OWNED
        if know.user_id != user_id:
            power = await KnowTeamedService.has_permission(know.code, user_id, KnowTeamedService.POWER_VIEW)

        # 返回结果
        return schema.KnowDetailVo(
            code=know.code,
            logo=await UrlUtil.to_absolute_url(know.logo),
            name=know.name,
            intro=know.intro,
            power=power,
            owned=True if know.user_id == user_id else False,
            vector_name=vector_name,
            vector_model=know.vector_model,
            splits_model=know.splits_model,
            is_disable=know.is_disable,
            create_time=TimeUtil.timestamp_to_date(know.create_time),
            update_time=TimeUtil.timestamp_to_date(know.update_time)
        )

    @classmethod
    async def add(cls, post: schema.KnowAddIn, user_id: int):
        """
        知识库创建。

        Args:
            post (schema.KnowAddIn): 知识库创建参数。
            user_id (int): 用户ID。

        Author:
            zero
        """
        try:
            # 验证剩余数
            surplus: int = await SubscribeLogic.get_surplus_quota(user_id, scene="know")
            if surplus <= 0:
                raise AppException("知识库创建额度不足!")

            # 验证模型
            await DevModelsModel.check_models(post.vector_model, ModelsEnum.TYPE_VECTOR)
            await DevModelsModel.check_models(post.splits_model, ModelsEnum.TYPE_CHAT)

            # 知识库Logo
            if not post.logo:
                post.logo = "static/images/default_know.png"

            # 违规审查
            review = CensorDriver.REVIEW_TEXTS
            values = post.name + "\n" + post.intro
            censor_info = await CensorDriver.text(user_id=int(user_id), text=values, scene=review)
            if censor_info is not None:
                raise AppException(f"知识库名称违规: {censor_info}")

            # 创建数据
            code: str = await ToolsUtil.make_rand_sn(AgentKnowModel, "code", 16, "string")
            await AgentKnowModel.create(
                code=code,
                user_id=user_id,
                logo=UrlUtil.to_relative_url(post.logo),
                name=post.name,
                intro=post.intro,
                vector_model=post.vector_model,
                splits_model=post.splits_model,
                create_time=int(time.time()),
                update_time=int(time.time())
            )
        except Exception as e:
            raise AppException(str(e))

    @classmethod
    async def edit(cls, post: schema.KnowEditIn, user_id: int):
        """
        知识库编辑。

        Args:
            post (schema.KnowAddIn): 知识库编辑参数。
            user_id (int): 用户ID。

        Author:
            zero
        """
        try:
            # 知识库校验
            know = await AgentKnowModel.filter(code=post.code, is_delete=0).get()

            # 权限的校验
            await KnowTeamedService.has_permission(know.code, user_id)

            # 验证模型
            await DevModelsModel.check_models(post.vector_model, ModelsEnum.TYPE_VECTOR)
            await DevModelsModel.check_models(post.splits_model, ModelsEnum.TYPE_CHAT)

            # 违规审查
            review = CensorDriver.REVIEW_TEXTS
            values = post.name + "\n" + post.intro
            censor_info = await CensorDriver.text(user_id=int(user_id), text=values, scene=review)
            if censor_info is not None:
                raise AppException(f"知识库信息违规: {str(censor_info)}")

            # 更新数据
            await AgentKnowModel.filter(code=post.code).update(
                logo=UrlUtil.to_relative_url(post.logo),
                name=post.name,
                intro=post.intro,
                vector_model=post.vector_model,
                splits_model=post.splits_model,
                update_time=int(time.time())
            )
        except Exception as e:
            raise AppException(str(e))

    @classmethod
    async def delete(cls, code: str, user_id: int):
        """
        知识库删除。

        Args:
            code (str): 知识库编号。
            user_id (int): 用户ID。

        Author:
            zero
        """
        # 知识库验证
        know = await AgentKnowModel.filter(code=code).get()

        # 权限的校验
        await KnowTeamedService.has_permission(know.code, user_id, KnowTeamedService.POWER_OWNED)

        async with in_transaction("mysql"):
            # 知识库删除
            await AgentKnowModel.filter(code=code).update(
                is_delete=1,
                delete_time=int(time.time())
            )
            # 关联表清理
            await CleanLogic.clean_robot(know.id)

    @classmethod
    async def change_vector(cls, code: str, user_id: int, vector_model: str):
        """
        变更向量模型。

        Args:
            code (str): 知识库编号。
            user_id (int): 用户ID。
            vector_model (str): 向量模型。

        Author:
            zero
        """
        # 知识库验证
        know = await AgentKnowModel.filter(code=code).get()

        # 权限的校验
        await KnowTeamedService.has_permission(know.code, user_id, KnowTeamedService.POWER_OWNED)

        # 向量模型校验
        try:
            models = await DevModelsModel.check_models(vector_model, ModelsEnum.TYPE_VECTOR)
        except Exception as e:
            raise AppException(str(e))

        # 验证训练数据
        wait_count = await (PgKnowledgeModel
                            .filter(know_id=know.id)
                            .filter(is_delete=0)
                            .filter(status__in=[AgentEnum.EMB_WAIT, AgentEnum.EMB_LINE, AgentEnum.EMB_ING])
                            .count())
        if wait_count > 0:
            raise AppException("知识库中存在『待训练』或『训练中』的数据，不可切换。")

        # 修改向量模型
        await AgentKnowModel.filter(code=code).update(
            vector_model=vector_model,
            update_time=int(time.time())
        )

        # 更新训练数据
        async with in_transaction("pgsql"):
            await (PgKnowledgeModel
                   .filter(know_id=know.id)
                   .filter(is_delete=0)
                   .update(
                        last_uid=user_id,
                        error="",
                        embedding=None,
                        dimension=0,
                        use_points=0,
                        use_tokens=0,
                        vector_model=models.get("model"),
                        vector_alias=models.get("alias"),
                        status=AgentEnum.EMB_WAIT,
                        update_time=time.time()
                   ))

            await (PgKnowledgeModel
                   .filter(know_id=know.id)
                   .filter(is_delete=1)
                   .update(
                        last_uid=user_id,
                        embedding=None,
                        dimension=0,
                        vector_model=models.get("model"),
                        vector_alias=models.get("alias"),
                        update_time=time.time()
                   ))
